import numpy as np
import pandas as pd
from sklearn.tree import DecisionTreeClassifier
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score, classification_report
from imblearn.over_sampling import SMOTE
from imblearn.combine import SMOTEENN
from sklearn.feature_selection import SelectKBest, chi2
from sklearn.ensemble import RandomForestClassifier
from xgboost import XGBClassifier
import sys
import io

# 设置标准输出的默认编码为UTF-8
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')

# 加载数据集
data = pd.read_excel("files/data.xlsx")  # 从Excel文件加载数据
df = pd.DataFrame(data)  # 转换为DataFrame格式

# 创建一个新的类别列，将 EMCI 和 LMCI 合并为 MCI
df['类别_MCI'] = df['类别'].replace({'EMCI': 'MCI', 'LMCI': 'MCI'})

# 定义任务列表
tasks = [
    ("AD", "CN"),  # 阿尔茨海默病（AD） vs 健康对照组（CN）
    ("MCI", "CN"),  # 轻度认知障碍（MCI） vs 健康对照组（CN）
    ("AD", "MCI"),  # 阿尔茨海默病（AD） vs 轻度认知障碍（MCI）
    ("AD", "EMCI"),  # 阿尔茨海默病（AD） vs 早期轻度认知障碍（EMCI）
    ("AD", "LMCI"),  # 阿尔茨海默病（AD） vs 晚期轻度认知障碍（LMCI）
    ("CN", "EMCI"),  # 健康对照组（CN） vs 早期轻度认知障碍（EMCI）
    ("CN", "LMCI"),  # 健康对照组（CN） vs 晚期轻度认知障碍（LMCI）
    ("EMCI", "LMCI"),  # 早期轻度认知障碍（EMCI） vs 晚期轻度认知障碍（LMCI）
    ("AD", "EMCI", "LMCI"),  # 三分类任务：AD vs EMCI vs LMCI
    ("CN", "MCI", "AD")  # 三分类任务：CN vs MCI vs AD
]

# 定义存储结果的列表
results = []

# 遍历每个任务
for task in tasks:
    # 如果任务是二分类，使用原来的方法
    if len(task) == 2:
        class1, class2 = task
        if class1 == 'MCI' or class2 == 'MCI':
            label_column = '类别_MCI'
        else:
            label_column = '类别'

        # 筛选属于当前任务的样本
        filtered_df = df[df[label_column].isin([class1, class2])]

        if filtered_df.shape[0] == 0:
            print(f"任务 {class1} vs {class2} 的数据集为空，跳过此任务。")
            continue

        # 提取特征和目标
        X = filtered_df.drop(['类别', '类别_MCI'], axis=1).select_dtypes(include=[np.number])
        y = filtered_df[label_column].map({class1: 0, class2: 1})

    # 如果任务是多分类，调整标签处理
    else:
        class1, class2, class3 = task
        label_column = '类别_MCI'
        filtered_df = df[df[label_column].isin([class1, class2, class3])]

        if filtered_df.shape[0] == 0:
            print(f"任务 {class1} vs {class2} vs {class3} 的数据集为空，跳过此任务。")
            continue

        # 提取特征和目标
        X = filtered_df.drop(['类别', '类别_MCI'], axis=1).select_dtypes(include=[np.number])
        y = filtered_df[label_column].map({class1: 0, class2: 1, class3: 2})  # 为多分类任务分配标签

    # 特征选择（使用卡方检验选择k个最佳特征）
    selector = SelectKBest(chi2, k=10)
    X = selector.fit_transform(X, y)

    # 处理不平衡数据（使用SMOTEENN综合采样方法）
    if len(np.unique(y)) == 1:
        X_resampled, y_resampled = X, y
    else:
        smote_enn = SMOTEENN(random_state=42)
        X_resampled, y_resampled = smote_enn.fit_resample(X, y)

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(
        X_resampled, y_resampled, test_size=0.20, random_state=15, stratify=y_resampled
    )

    # 数据标准化
    scaler = StandardScaler()
    scaled_X_train = scaler.fit_transform(X_train)
    scaled_X_test = scaler.transform(X_test)

    # 定义超参数网格
    param_grid = {
        'n_estimators': [100, 250],
        'max_depth': [3, 5, 7],
        'min_samples_split': [10, 20],
        'min_samples_leaf': [10, 20],
        'class_weight': [None, 'balanced']
    }

    # 进行超参数搜索（使用随机森林分类器）
    grid_search = GridSearchCV(
        RandomForestClassifier(random_state=40),
        param_grid,
        cv=10,
        scoring='accuracy',
        n_jobs=-1
    )
    grid_search.fit(scaled_X_train, y_train)

    # 获取最佳模型
    best_model = grid_search.best_estimator_

    # 对测试集进行预测
    y_pred = best_model.predict(scaled_X_test)

    # 计算准确率和分类报告
    accuracy = accuracy_score(y_test, y_pred)
    report = classification_report(y_test, y_pred, output_dict=True)

    # 存储结果
    results.append({
        '任务名称': f'{task}',  # 当前任务名称
        '最佳参数': grid_search.best_params_,
        '准确率': accuracy,
        '分类报告': report
    })

# 打印所有任务的结果
for result in results:
    print(f"任务名称: {result['任务名称']}")
    print(f"最佳参数: {result['最佳参数']}")
    print(f"准确率: {result['准确率']:.2%}")
    print("分类报告:")
    report_df = pd.DataFrame(result['分类报告']).T
    report_df.columns = ["精确率", "召回率", "F1分数", "支持样本数"]
    print(report_df.to_string())
    print("\n")